// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: xguo

#include "coin2/strategy/umm/umm.h"

#include "coin2/base/proto_util.h"
#include "coin2/base/string_util.h"
#include "coin2/exchange/base/log/latency_recorder.h"

namespace {

std::string GetKey(const IProduct& product) {
  // get base symbol
  return product.base().symbol();
}

uint64_t GetProductId(const IProduct& product) {
  return std::hash<std::string>()(
      product.absolute_norm() + "/" + product.exchange_name() + "/" + product.market_name());
}

double GetSign(OrderEvent event) {
  switch (event.trade_side()) {
    case TradeSide::TRADE_BUY_SIDE: {
      return 1.;
    }
    case TradeSide::TRADE_SELL_SIDE: {
      return -1.;
    }
    default: {
      NOTREACHED() << "Unknown order side";
      return 0.;
    }
  }
  return 0.;
}

coin::proto::LatencyProto::LatencyTag GetFeedLatencyTag(const FeedUpdate& upd) {
  if (upd.is_book())
    return coin::proto::LatencyProto::BOOK_FEED_LATENCY;
  else if (upd.is_trade())
    return coin::proto::LatencyProto::TRADE_FEED_LATENCY;
  else
    return coin::proto::LatencyProto::FEED_LATENCY;
}

coin::proto::AccountRequestProto GetAccRequest(const FeedUpdate& upd) {
  const MarketExchangeApi& mea = upd.mea();
  coin::proto::AccountRequestProto acc_request;
  acc_request.set_market_type(MarketType_Name(mea.market));
  acc_request.set_exchange(ExchangeType_Name(mea.exchange));
  acc_request.set_api_version(mea.api);
  return acc_request;
}

}  // namespace

//------------------------------------
// UmmSubStrategy
//------------------------------------
UmmSubStrategy::UmmSubStrategy(const UmmSubStrategyConfig& config, int64_t ts)
    : start_time_(ts), config_(config) {
  // products
  for (const auto& s : config.symbols()) {
    products_.push_back(CreateProductFromUniqueString(s, start_time_));
    LOG(INFO) << "add product: " << products_.back()->absolute_norm();
  }
  CHECK_LE(products_.size(), 16);
  CHECK(!products_.empty());

  pi_ = &(product_cache_.holder(*trade_product()).product_info());
  is_trading_futures_ = trade_product()->market() == MarketType::Futures;

  std::string env_dryrun = getenv("UMM_DRY_RUN") ? getenv("UMM_DRY_RUN") : "";
  dry_run_ = config_.dry_run() || (env_dryrun == "1");

  LOG(INFO) << "Create UmmSubStrategy: \n"
            << "is_trading_futures: " << is_trading_futures_ << "\n"
            << "dry_run: " << dry_run_;
}

void UmmSubStrategy::Init(Driver* driver) {
  constexpr int64_t T = 1'000'000'000LL;
  driver_ = driver;
  os_ = driver_->order()->GetUniqueSystem(mea());
  auto get_product_holder = [this](const coin2::exchange::base::symbology::IProduct& product, int64_t timestamp) {
      return driver_->product_holder_manager()->GetProductHolder(product, timestamp);
  };

  // create executor
  auto* pass_config = config_.mutable_pass_executor_config();
  CHECK(pass_config->has_order_update_period_sec());
  CHECK(pass_config->has_max_posting_period_sec());
  CHECK(pass_config->has_price_pull_bp());
  CHECK(pass_config->has_price_push_bp());

  pass_config->set_order_update_period(pass_config->order_update_period_sec() * T);
  pass_config->set_max_posting_period(pass_config->max_posting_period_sec() * T);
  LOG(INFO) << "passive executor config: " << config_.pass_executor_config().DebugString();
  executor_ =
      std::make_unique<PassiveOrderExecutor>(config_.pass_executor_config(), os(), start_time_, nullptr, get_product_holder);

  // create pricer
  auto pconfig = config_.pricer_config();
  xguo::UmmPricerConfig config;

  std::transform(products_.begin(), products_.end(), config.product_ids, [](const auto& p) {
    return GetProductId(*p);
  });

  snprintf(config.name, sizeof(config.name), "%s", pconfig.name().data());
  config.num_products = products_.size();
  config.maker_fee_rate = pconfig.maker_fee_bp() * 1e-4;
  config.taker_fee_rate = pconfig.taker_fee_bp() * 1e-4;
  config.walk_step = pconfig.walk_step_bp() * 1e-4;

  config.ohlc_window_size = pconfig.ohlc_window_size_sec() * T;
  config.rate_window_size = pconfig.rate_window_size_sec() * T;
  config.max_nofill_window_size = pconfig.max_nofill_window_size_sec() * T;

  config.risk_threshold = pconfig.risk_threshold();
  config.take_profit_rate = pconfig.take_profit_bp() * 1e-4;
  config.stop_loss_rate = pconfig.stop_loss_bp() * 1e-4;
  config.force_fire_cooldown_time = pconfig.force_fire_cooldown_sec() * T;
  config.use_agg = pconfig.use_agg();
  config.force_on_level = pconfig.force_on_level();
  config.max_level = pconfig.max_level();
  config.min_level = pconfig.min_level();
  config.max_pos = pconfig.max_pos();
  config.min_pos = pconfig.min_pos();
  config.reserve = (config.min_pos + config.max_pos) * 0.5;

  CHECK_GT(config.take_profit_rate, 0);
  CHECK_LT(config.stop_loss_rate, 0);

  config.reserve = (config.max_pos + config.min_pos) * 0.5;
  config.lot_size = pconfig.lot_size();
  config.order_update_period_sec = config_.pass_executor_config().order_update_period_sec();

  pricer_ = std::make_unique<xguo::UmmPricer>(config);
}

void UmmSubStrategy::onBookFeed(const FeedUpdate& upd) {
  auto book = upd.GetBookBuilder();
  auto ask0 = book->Ask0();
  auto bid0 = book->Bid0();

  xguo::Bbo bbo;
  bbo.product_id = GetProductId(upd.product());
  bbo.ask0 = ask0->price;
  bbo.bid0 = bid0->price;
  bbo.midp = (bbo.ask0 + bbo.bid0) * 0.5;
  bbo.timestamp = timestamp_;

  // Do some simple check.
  CHECK_GT(bbo.timestamp, 10000) << bbo.timestamp;
  if (bbo.bid0 == 0.0 || bbo.ask0 == 0.0) {
    LOG(ERROR) << "book error! ";
    return;
  }

  pricer_->UpdateByBbo(bbo);

  if (!is_ready()) {
    LOG(INFO) << "Pricer/OG not ready! " << pricer_->name() << ", " << pricer_->num_feeds() << ", "
              << "pricer: " << pricer_->is_ready() << ", "
              << "OG: " << os()->is_ready();
    return;
  }
  LOG_FIRST_N(INFO, 1) << "pricer_og_ready: " << pricer_->name();

  const auto& order_data = pricer_->CalculateOrder();
  if (!order_data.success) {
    return;
  }

  /*
  order_data.CheckPrice();
  if (timestamp() >  x * T && timestamp() < (x+3600) * T) {
    LOG(ERROR) << "+++++++++++++++++++++++++";
    LOG(ERROR) << order_data.DebugString();
    LOG(ERROR) << "+++++++++++++++++++++++++";
    LOG(ERROR) << pricer_->DebugString();
    LOG(ERROR) << "+++++++++++++++++++++++++";
  }
  */

  const auto& tick_price = pi_->tick_price();
  const auto& tick_qty = pi_->tick_qty();

  // DLOG(ERROR) << order_data.tag;
  if (order_data.type == xguo::OrderData::PASSIVE) {
    double sell_price = 0;
    double sell_qty = 0;

    if (order_data.pass_sell_price > 0) {
      sell_price = tick_price.GetNextPrice(order_data.pass_sell_price);
      sell_qty = tick_qty.Round(order_data.pass_sell_qty);
    }

    double buy_price = 0;
    double buy_qty = 0;
    if (order_data.pass_buy_price > 0) {
      buy_price = tick_price.GetPrevPrice(order_data.pass_buy_price);
      buy_qty = tick_qty.Round(order_data.pass_buy_qty);
    }

    if (risk_manager_->num_of_rest_errors() > 50) {
      LOG(ERROR) << "num_of_rest_errors: " << risk_manager_->num_of_rest_errors()
                 << "window_size: " << risk_manager_->rest_stat_window_size();
      return;
    }

    if (!dry_run_) {
      executor_->SetMmTag(order_data.tag);
      executor_->ManageMmOrders(timestamp_, sell_price, buy_price, sell_qty, buy_qty, false, false);
    }
    return;
  }

  if (order_data.type == xguo::OrderData::AGGRESSIVE_BUY) {
    LOG(INFO) << "FireAggBuy";
    if (!dry_run_) {
      executor_->ForceFireBuyOrder(
          timestamp_,
          tick_price.GetPrevPrice(order_data.agg_buy_price),
          tick_qty.Round(order_data.agg_buy_qty),
          order_data.tag);
    }
    return;
  }

  if (order_data.type == xguo::OrderData::AGGRESSIVE_SELL) {
    LOG(INFO) << "FireAggSell";
    if (!dry_run_) {
      executor_->ForceFireSellOrder(
          timestamp_,
          tick_price.GetNextPrice(order_data.agg_sell_price),
          tick_qty.Round(order_data.agg_sell_qty),
          order_data.tag);
    }
    return;
  }

  NOTREACHED() << "Unknown type!";
}

void UmmSubStrategy::onTradeFeed(const FeedUpdate& upd) { (void)upd; }

void UmmSubStrategy::onAccountOrder(const OrderUpdate& upd) {
  using FillType = coin::proto::FillType;
  auto event = upd.product_order_info().event();
  if (event.type() != OrderEvent::ORDER_FILLED) {
    return;
  }

  xguo::Fill fill;
  fill.product_id = GetProductId(upd.product());
  fill.timestamp = timestamp_;
  fill.fill_price = event.fill_price();
  fill.fill_qty = GetSign(event) * event.fill_qty();

  const double amt = event.fill_price() * event.fill_qty();
  LOG(INFO) << "fill_received: " << event.DebugString();
  if (event.fill_type() == FillType::MAKER_FILL_TYPE) {
    fill.fee = std::fabs(pricer_->config().maker_fee_rate * amt);
  } else if (event.fill_type() == FillType::TAKER_FILL_TYPE) {
    fill.fee = std::fabs(pricer_->config().taker_fee_rate * amt);
  } else {
    double rate = 0.5 * (pricer_->config().maker_fee_rate + pricer_->config().taker_fee_rate);
    fill.fee = std::fabs(rate * amt);
  }

  pricer_->UpdateByFill(fill);
  LOG(INFO) << pricer_->DebugString();
}

void UmmSubStrategy::onAccountInfo(const OrderUpdate& upd) {
  auto& account_info = upd.account_info();
  if (!account_info.is_ready()) {
    return;
  }

  if (is_trading_futures_ && account_info.is_position_update()) {
    const auto& prod = *trade_product();
    if (account_info.HasPosition(prod)) {
      double pos = account_info.GetPosition(prod).net_position();
      pricer_->UpdateByPos(pos);
      executor_->onAccountInfo(upd);
      num_pos_update_++;
      LOG_FIRST_N(INFO, 3) << "umm_position: " << trade_product()->absolute_norm() << ", " << pos;
    }
  } else {
    const auto& base = trade_product()->base();
    if (account_info.HasBalance(base)) {
      double bal = account_info.GetBalance(trade_product()->base()).total();
      pricer_->UpdateByPos(bal);
      executor_->onAccountInfo(upd);
      num_pos_update_++;
      LOG_FIRST_N(INFO, 3) << "umm_balance: " << trade_product()->absolute_norm() << ", " << bal;
    }
  }
}

//------------------------------------
// UmmStrategy
//------------------------------------
UmmStrategy::UmmStrategy(const std::string& config_file, std::optional<int64_t> ts) {
  std::string content = ReadFileContent(config_file);
  LOG(INFO) << "strategy config json: \n" << content;

  CHECK(JsonStringToProto(content, &config_)) << content;
  LOG(INFO) << "strategy config proto: \n" << config_.DebugString();

  RiskManagerConfig risk_config;
  risk_config.set_rest_error_window_size_sec(config_.rest_error_window_size_sec());
  risk_manager_ = std::make_unique<RiskManager>(risk_config);

  start_time_ = ts.value_or(GetCurrentTimestamp());
  for (auto& [key, subconfig] : config_.subconfig()) {
    auto strat = std::make_unique<UmmSubStrategy>(subconfig, start_time_);
    strat->SetRiskManager(risk_manager_.get());
    strats_.emplace(key, std::move(strat));
  }
}

void UmmStrategy::UpdateSimConfig(SimConfig* config) {
  LOG(INFO) << "sim_config: " << config->DebugString();
  for (const auto& p : reserve_map_) {
    config->add_reserve_symbols(p.first->absolute_norm());
    config->add_reserve_positions(p.second);
    config->set_start_timestamp(start_time_);
  }
}

void UmmStrategy::Init(Driver* driver) {
  driver_ = driver;
  os_ = driver_->order()->GetUniqueSystem(mea());
  for (auto&& strat : strats_) {
    strat.second->Init(driver);
    reserve_map_[strat.second->trade_product()] = strat.second->reserve();
  }

  auto interval = driver->strategy()->config().pnl_report_period_ns();
  if (interval > 0) {
    pnl_monitor_ =
        std::make_unique<PnlMonitor>(driver->live_executor(), driver->archive_executor(), interval);
    pnl_monitor_->StartAsync();
  }

  LOG(INFO) << "Set Latency Recorder";
  std::unordered_map<std::string, double> map;

  // TODO(xguo) handle position
  for (auto& [p, v] : reserve_map_) {
    LOG(INFO) << "reserve: " << p->absolute_norm() << ", " << v;
    bool is_trading_futures = p->market() == MarketType::Futures;
    if (is_trading_futures) {
      map[p->absolute_norm()] = v;
    } else {
      map[p->base().currency()] = v;
    }
  }
  strat_reporter_.reset(new StrategyReporter(driver->strategy()->config(), map));
  for (auto& strat : strats_) {
    strat_reporter_->RegisterProduct(*strat.second->trade_product());
  }

  os()->gateway()->set_latency_recorder(
      strat_reporter_->GetMutableStrategyLogger()->mutable_latency_recorder());
}

void SetFeedTimeToLatency(
    const FeedUpdate& upd,
    coin::proto::LatencyProto::LatencyTag tag,
    const coin::proto::AccountRequestProto& acc_request,
    bool mark_end,
    coin2::exchange::base::strategy_util::LatencyRecorder* latency_recorder) {
  if (!latency_recorder) {
    LOG(ERROR) << "doesn't have latency recorder";
    return;
  }
  latency_recorder->Begin(tag);
  auto latency_context_id = latency_recorder->GetCurrentContextId();
  auto feed_publish_point = coin::proto::LifeOfSignal::RAW_FEED_EXCHANGE_PUBLISHED;
  if (upd.is_book()) {
    feed_publish_point = coin::proto::LifeOfSignal::RAW_BOOK_FEED_EXCHANGE_PUBLISHED;
  } else if (upd.is_trade()) {
    feed_publish_point = coin::proto::LifeOfSignal::RAW_TRADE_FEED_EXCHANGE_PUBLISHED;
  }
  latency_recorder->SetProtoAndTime(
      latency_context_id,
      feed_publish_point,
      upd.ts().exchange_publish_timestamp(),
      acc_request);
  latency_recorder->SetProtoAndTime(
      latency_context_id,
      coin::proto::LifeOfSignal::RAW_FEED_RECEIVED,
      upd.ts().raw_rx_timestamp(),
      acc_request);
  latency_recorder->SetProtoAndTime(
      latency_context_id,
      coin::proto::LifeOfSignal::RAW_FEED_PARSE_STARTED,
      upd.ts().main_rx_timestamp(),
      acc_request);
  if (upd.is_book()) {
    latency_recorder->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::ON_BOOK_FEED,
        upd.ts().main_tx_timestamp(),
        acc_request);
  } else if (upd.is_trade()) {
    latency_recorder->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::ON_TRADE_FEED,
        upd.ts().main_tx_timestamp(),
        acc_request);
  }
  if (mark_end) {
    latency_recorder->End(latency_context_id);
  }
}

void UmmStrategy::onBookFeed(const FeedUpdate& upd) {
  const int64_t ts = upd.ts().exchange_publish_timestamp();
  UpdateTimestamp(ts);

  if (timestamp_ > 10000) {
    const auto key = GetKey(upd.product());
    strats_[key]->onBookFeed(upd);
  }

  auto* latency_recorder = strat_reporter_->GetMutableStrategyLogger()->mutable_latency_recorder();
  auto acc_request = GetAccRequest(upd);
  auto feed_latency_tag = GetFeedLatencyTag(upd);
  SetFeedTimeToLatency(upd, feed_latency_tag, acc_request, true, latency_recorder);
  SetFeedTimeToLatency(
      upd,
      coin::proto::LatencyProto::ORDER_LATENCY,
      acc_request,
      false,
      latency_recorder);
}

void UmmStrategy::onTradeFeed(const FeedUpdate& upd) {
  const int64_t ts = upd.ts().exchange_publish_timestamp();
  UpdateTimestamp(ts);

  if (timestamp_ > 10000) {
    const auto key = GetKey(upd.product());
    strats_[key]->onTradeFeed(upd);
  }
}

void UmmStrategy::onAccountOrder(const OrderUpdate& upd) {
  const auto key = GetKey(upd.product());
  strats_[key]->onAccountOrder(upd);
}

void UmmStrategy::onAccountInfo(const OrderUpdate& upd) {
  auto& account_info = upd.account_info();

  if (!account_info.is_ready()) {
    return;
  }

  for (auto& strat : strats_) {
    strat.second->onAccountInfo(upd);
  }

  if (strat_reporter_) {
    for (auto& strat : strats_) {
      if (strat.second->is_ready()) {
        auto* prod = strat.second->trade_product();
        double price = strat.second->trade_price();
        CHECK_GT(price, 0);
        strat_reporter_->UpdateMidp(*prod, price);
      }
    }
    strat_reporter_->onAccountInfo(upd);
  }

  constexpr int64_t timeout = 10'000'000'000LL;
  if (timestamp_ - last_print_ts_ > timeout) {
    LOG(INFO) << os()->gateway()->order_manager().DebugString();
    last_print_ts_ = timestamp_;
  }
}

void UmmStrategy::onOrderLog(const coin::proto::OrderGatewayLog& order_log) {
  if (strat_reporter_) {
    auto logger = strat_reporter_->GetMutableStrategyLogger();
    logger->WriteOrderLog(order_log, true, GetCurrentTimestamp());
  }
}

void UmmStrategy::onOrderSubsystemError(const OrderUpdate& upd) {
  LOG(INFO) << "OrderSubsystemError: " << coin::proto::OrderUpdateType_Name(upd.update_type());
  LOG(INFO) << "OrderSubsystemError: " << upd.timestamp();
  risk_manager_->HandleOrderUpdate(upd);
}

void UmmStrategy::onCleanupAux(Driver* driver, int num_retries) {
  auto ts = GetCurrentTimestamp();
  int num_orders = 0;
  for (auto& strat : strats_) {
    auto* executor = strat.second->Executor();
    if (executor) {
      executor->SetEnableOrderSubmission(false);
      num_orders += executor->TryCancelAllWorkingOrders(ts);
    }
  }

  if (num_orders == 0) {
    LOG(INFO) << "Cleared all pending orders. " << num_retries;
    driver->PostTask([driver](){driver->Exit();});
  } else if (num_retries > 10) {
    LOG(WARNING) << "There may be some pendings orders left open.";
    driver->PostTask([driver](){driver->Exit();});
  }
}

void UmmStrategy::onCleanup(Driver* driver) {
  driver->SetImmediateExitOnSignal(false);
  driver->PostPeriodicTask([this, driver](){ this->onCleanupAux(driver, 0); }, 500'000'000LL);
}